package com.wfh.aiagent.chatmemory;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import org.objenesis.strategy.StdInstantiatorStrategy;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * 基于文件的聊天记忆实现，使用 Kryo 序列化消息列表。
 *
 * @Author FengHuan Wang
 * @Date 2025/5/25 12:35
 * @Version 1.0
 */
public class FileBasedMemory implements ChatMemory {

    private final String BASE_DIR;

    private static final Kryo kryo = new Kryo();

    static {
        kryo.setRegistrationRequired(false);
        kryo.register(org.springframework.ai.chat.messages.UserMessage.class);
        kryo.register(org.springframework.ai.chat.messages.AssistantMessage.class);
        kryo.register(org.springframework.ai.chat.messages.SystemMessage.class);
        kryo.register(ArrayList.class);
        kryo.setInstantiatorStrategy(new StdInstantiatorStrategy());
    }

    public FileBasedMemory(String baseDir) {
        this.BASE_DIR = baseDir;
        File dir = new File(baseDir);
        if (!dir.exists()) {
            dir.mkdirs();
        }
    }

    // ✅ 正确实现接口方法：返回完整消息列表
    public List<Message> get(String conversationId) {
        return getOrCreateConversation(conversationId);
    }

    // ✅ 正确实现接口方法：批量添加消息
    @Override
    public void add(String conversationId, List<Message> messages) {
        if (messages == null || messages.isEmpty()) {
            return;
        }
        List<Message> existing = getOrCreateConversation(conversationId);
        existing.addAll(messages);
        saveConversation(conversationId, existing);
    }


    @Override
    public void add(String conversationId, Message message) {
        if (message != null) {
            add(conversationId, List.of(message));
        }
    }

    // ✅ 自定义扩展方法（非接口要求）
    @Override
    public List<Message> get(String conversationId, int lastN) {
        List<Message> all = getOrCreateConversation(conversationId);
        int fromIndex = Math.max(0, all.size() - lastN);
        return new ArrayList<>(all.subList(fromIndex, all.size()));
    }

    // ✅ 实现 clear 方法
    @Override
    public void clear(String conversationId) {
        File file = getConversationFile(conversationId);
        if (file.exists()) {
            if (!file.delete()) {
                System.err.println("Failed to delete conversation file: " + file.getAbsolutePath());
            }
        }
    }

    /**
     * 从文件加载或创建空消息列表
     */
    private List<Message> getOrCreateConversation(String conversationId) {
        File file = getConversationFile(conversationId);
        if (!file.exists()) {
            return new ArrayList<>();
        }

        try (Input input = new Input(new FileInputStream(file))) {
            @SuppressWarnings("unchecked")
            List<Message> messages = kryo.readObject(input, ArrayList.class);
            if (messages == null) {
                return new ArrayList<>();
            }
            // 过滤 null（防御性编程）
            List<Message> validMessages = new ArrayList<>();
            for (Message msg : messages) {
                if (msg != null) {
                    validMessages.add(msg);
                }
            }
            return validMessages;
        } catch (Exception e) {
            System.err.println("Failed to load conversation: " + conversationId + " from " + file.getAbsolutePath());
            e.printStackTrace();
            return new ArrayList<>(); // 容错：返回空列表
        }
    }

    /**
     * 将消息列表保存到文件
     */
    private void saveConversation(String conversationId, List<Message> messages) {
        File file = getConversationFile(conversationId);
        try (Output output = new Output(new FileOutputStream(file))) {
            kryo.writeObject(output, messages);
        } catch (IOException e) {
            System.err.println("Failed to save conversation: " + conversationId + " to " + file.getAbsolutePath());
            e.printStackTrace();
        }
    }

    /**
     * 获取会话对应的文件路径
     */
    private File getConversationFile(String conversationId) {
        return new File(BASE_DIR, conversationId + ".bin");
    }
}